﻿/*
* Copyright (c) 2012 Kyle Hayward
* 
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to deal
* in the Software without restriction, including without limitation the rights
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
* copies of the Software, and to permit persons to whom the Software is
* furnished to do so, subject to the following conditions:
* 
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
* 
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
* THE SOFTWARE.
*/
using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Data;
using System.Drawing;
using System.Linq;
using System.Text;
using System.IO;
using System.Windows.Forms;
using System.Security;
using System.Runtime.InteropServices;

using Tao.OpenGl;
using Tao.FreeGlut;
using Tao.Platform.Windows;

using OptixDotNet;
using OptixDotNet.Utils;
using OptixDotNet.Math;

namespace OptixDotNetUtils
{
    public partial class OptixForm : Form
    {
        [SuppressUnmanagedCodeSecurity]
        [DllImport( "user32.dll", CharSet = CharSet.Auto )]
        [return: MarshalAs( UnmanagedType.Bool )]
        public static extern bool PeekMessage( out Message msg, IntPtr hWnd, uint messageFilterMin, uint messageFilterMax, uint flags );

        #region Public Fields

        public ICamera Camera = null;

        public Context Context = null;
        public OptixDotNet.Buffer OutputBuffer = null;

        public Vector3 BackBufferColor;

        public new int Width { get { return GlControl.Width; } }
        public new int Height { get { return GlControl.Height; } }

        public int RaysTracedPerFrame;

        public bool UsePBO = true;
        public bool UseSRGB = false;
        public bool DrawUI = true;

        #endregion

        #region Private Fields

        //stats
        private int mNumFrames = 0;
        private int mNumRaysTraced = 0;
        private float mTimeElapsed = 0.0f;
        private float mTimeDelta = 0.0f;
        private string mText = "";

        //ogl
        private uint mPBOTexID = 0;

        //mouse
        private MouseEventArgs mPrevMouse;

        #endregion

        public OptixForm()
        {
            InitializeComponent();

            this.SetStyle( ControlStyles.AllPaintingInWmPaint | ControlStyles.UserPaint, true );

            GlControl.KeyUp += new KeyEventHandler( GLControl_KeyUp );
            GlControl.KeyDown += new KeyEventHandler( GlControl_KeyDown );
            GlControl.MouseMove += new MouseEventHandler( GLControl_MouseMove );
            GlControl.MouseDown += new MouseEventHandler( GLControl_MouseDown );
            GlControl.MouseUp += new MouseEventHandler( GLControl_MouseUp );

            GlControl.InitializeContexts();
            Glut.glutInit();
            Glut.glutInitDisplayMode( Glut.GLUT_RGB | Glut.GLUT_DOUBLE );

            Time.InitTime();

            mPrevMouse = new MouseEventArgs( MouseButtons.None, 0, 0, 0, 0 );

            BackBufferColor = new Vector3( 100 / 255.0f, 149 / 255.0f, 237 / 255.0f );
        }

        private bool AppStillIdle()
        {
            Message msg;
            return !PeekMessage( out msg, IntPtr.Zero, 0, 0, 0 );
        }

        public void OnIdle( object sender, EventArgs e )
        {
            while ( AppStillIdle() )
            {
                double start = Time.GetTimeInSecsDouble();
                RayTrace();
                double stop = Time.GetTimeInSecsDouble();

                mTimeElapsed += (float)( stop - start );
                mNumFrames++;
                mNumRaysTraced += RaysTracedPerFrame;

                if ( mTimeElapsed >= 0.25f )
                {
                    float raysPerSec = mNumRaysTraced / mTimeElapsed;
                    raysPerSec /= 1000.0f * 1000.0f;

                    mText = string.Format( "ray-trace: {0:.00} ms : {1:.00} MRays/sec", mTimeElapsed / mNumFrames * 1000.0f, raysPerSec );
                    mTimeElapsed = 0.0f;
                    mNumFrames = 0;
                    mNumRaysTraced = 0;
                }

                // Initialize state
                Gl.glMatrixMode( Gl.GL_PROJECTION );
                Gl.glLoadIdentity();
                Gl.glOrtho( 0, 1, 0, 1, -1, 1 );
                Gl.glMatrixMode( Gl.GL_MODELVIEW );
                Gl.glLoadIdentity();
                Gl.glViewport( 0, 0, Width, Height );

                Gl.glClearColor( BackBufferColor.X, BackBufferColor.Y, BackBufferColor.Z, 1 );
                Gl.glClear( Gl.GL_COLOR_BUFFER_BIT );

                Display();

                DrawText( 15, 15, mText );

                System.Threading.Thread.Sleep( 10 );
                GlControl.Invalidate();

                stop = Time.GetTimeInSecsDouble();
                mTimeDelta = (float)( stop - start );
            }
        }

        protected override void OnLoad( EventArgs e )
        {
            base.OnLoad( e );

            try
            {
                Initialize();

                if( UseSRGB )
                    Gl.glEnable( Gl.GL_FRAMEBUFFER_SRGB_EXT );
            }
            catch ( Exception ex )
            {
                string text = string.Format( "{0}\n\n{1}", ex.Message, ex.StackTrace );
                MessageBox.Show( text, "Exception during Optix initialization!", MessageBoxButtons.AbortRetryIgnore );
                Application.Exit();
                Environment.Exit( 0 );
            }
        }

        protected virtual void Initialize()
        {

        }

        protected virtual void Shutdown()
        {
            if ( Context != null )
            {
                Context.Destroy();
                Context = null;
            }
        }

        protected virtual void RayTrace()
        {

        }

        protected virtual void Display()
        {
            if ( UsePBO )
            {
                DrawPBO();
            }
            else
            {
                BufferStream stream = OutputBuffer.Map();

                int type = 0;
                int format = 0;
                switch ( OutputBuffer.Format )
                {
                    case Format.UByte4:
                        format = Gl.GL_BGRA;
                        type = Gl.GL_UNSIGNED_BYTE;
                        break;
                    case Format.Float4:
                        format = Gl.GL_RGBA;
                        type = Gl.GL_FLOAT;
                        break;
                    case Format.Float3:
                        format = Gl.GL_RGB;
                        type = Gl.GL_FLOAT;
                        break;
                    case Format.Float:
                        format = Gl.GL_LUMINANCE;
                        type = Gl.GL_FLOAT;
                        break;
                    default:
                        throw new Exception( "Unsupported format for glDrawPixels" );
                }

                Gl.glDrawPixels( Width, Height, format, type, stream.DataPointer );

                OutputBuffer.Unmap();
            }
        }

        protected virtual void DrawText( int x, int y, string text )
        {
            if ( !DrawUI )
                return;

            Gl.glMatrixMode( Gl.GL_PROJECTION );
            Gl.glLoadIdentity();
            Gl.glOrtho( 0, Width, 0, Height, -1, +1 );

            // Save state
            Gl.glPushAttrib( Gl.GL_CURRENT_BIT | Gl.GL_ENABLE_BIT );

            Gl.glDisable( Gl.GL_TEXTURE_2D );
            Gl.glDisable( Gl.GL_LIGHTING );
            Gl.glDisable( Gl.GL_DEPTH_TEST );

            Gl.glColor3f( .1f, .1f, .1f ); // drop shadow color

            // Shift shadow one pixel to the lower right.
            Gl.glRasterPos2i( x + 1, y - 1 );
            foreach ( char c in text )
                Glut.glutBitmapCharacter( Glut.GLUT_BITMAP_8_BY_13, (int)c );

            Gl.glColor3f( 1.0f, 1.0f, 1.0f );        // main text
            Gl.glRasterPos2i( x, y );
            foreach ( char c in text )
                Glut.glutBitmapCharacter( Glut.GLUT_BITMAP_8_BY_13, (int)c );

            // Restore state
            Gl.glPopAttrib();
        }

        protected virtual void CameraUpdate()
        {

        }

        protected virtual bool LoadAccelData( OptixOBJLoader model, string sourcePath, string cachedAccelPath )
        {
            if ( !File.Exists( sourcePath ) || !File.Exists( cachedAccelPath ) )
                return false;

            if ( File.GetLastWriteTime( sourcePath ) > File.GetLastWriteTime( cachedAccelPath ) )
                return false;

            return model.GeoGroup.Acceleration.Load( cachedAccelPath );
        }

        protected void CreateOutputBuffer( Format format )
        {
            if ( UsePBO )
            {
                //Allocate first the memory for the gl buffer, then attach it to OptiX.
                uint pbo = CreatePBO( format );

                OGLBufferDesc glDesc = new OGLBufferDesc()
                {
                    Width = (uint)Width,
                    Height = (uint)Height,
                    Format = format,
                    Type = BufferType.Output,
                    Resource = pbo
                };
                OutputBuffer = new OGLBuffer( Context, glDesc );
            }
            else
            {
                BufferDesc desc = new BufferDesc() 
                { 
                    Width = (uint)Width, 
                    Height = (uint)Height, 
                    Format = format, 
                    Type = BufferType.Output 
                };
                OutputBuffer = new OptixDotNet.Buffer( Context, desc );
            }
        }

        protected uint CreatePBO( Format format )
        {
            uint pbo = 0;
            Gl.glGenBuffers( 1, out pbo );
            Gl.glBindBuffer( Gl.GL_ARRAY_BUFFER, pbo );

            int size = (int)Utils.GetFormatSize( format );
            Gl.glBufferData( Gl.GL_ARRAY_BUFFER, new IntPtr( size * Width * Height ), IntPtr.Zero, Gl.GL_STREAM_DRAW );
            Gl.glBindBuffer( Gl.GL_ARRAY_BUFFER, 0 );

            return pbo;
        }

        protected void DrawPBO()
        {
            OGLBuffer buffer = OutputBuffer as OGLBuffer;

            if ( buffer == null )
                return;

            if ( mPBOTexID == 0 )
            {
                Gl.glGenTextures( 1, out mPBOTexID );
                Gl.glBindTexture( Gl.GL_TEXTURE_2D, mPBOTexID );

                // Change these to GL_LINEAR for super- or sub-sampling
                Gl.glTexParameteri( Gl.GL_TEXTURE_2D, Gl.GL_TEXTURE_MAG_FILTER, Gl.GL_NEAREST );
                Gl.glTexParameteri( Gl.GL_TEXTURE_2D, Gl.GL_TEXTURE_MIN_FILTER, Gl.GL_NEAREST );

                // GL_CLAMP_TO_EDGE for linear filtering, not relevant for nearest.
                Gl.glTexParameteri( Gl.GL_TEXTURE_2D, Gl.GL_TEXTURE_WRAP_S, Gl.GL_CLAMP_TO_EDGE );
                Gl.glTexParameteri( Gl.GL_TEXTURE_2D, Gl.GL_TEXTURE_WRAP_T, Gl.GL_CLAMP_TO_EDGE );

                Gl.glBindTexture( Gl.GL_TEXTURE_2D, 0 );
            }

            Gl.glBindTexture( Gl.GL_TEXTURE_2D, mPBOTexID );

            // send pbo to texture
            Gl.glBindBuffer( Gl.GL_PIXEL_UNPACK_BUFFER, buffer.Resource );

            uint elementSize = (uint)buffer.ElementSize;
            if ( ( elementSize % 8 ) == 0 ) Gl.glPixelStorei( Gl.GL_UNPACK_ALIGNMENT, 8 );
            else if ( ( elementSize % 4 ) == 0 ) Gl.glPixelStorei( Gl.GL_UNPACK_ALIGNMENT, 4 );
            else if ( ( elementSize % 2 ) == 0 ) Gl.glPixelStorei( Gl.GL_UNPACK_ALIGNMENT, 2 );
            else Gl.glPixelStorei( Gl.GL_UNPACK_ALIGNMENT, 1 );

            switch ( buffer.Format )
            {
                case Format.UByte4:
                    Gl.glTexImage2D( Gl.GL_TEXTURE_2D, 0, Gl.GL_RGBA8, Width, Height, 0, Gl.GL_BGRA, Gl.GL_UNSIGNED_BYTE, IntPtr.Zero );
                    break;
                case Format.Float4:
                    Gl.glTexImage2D( Gl.GL_TEXTURE_2D, 0, Gl.GL_RGBA32F_ARB, Width, Height, 0, Gl.GL_RGBA, Gl.GL_FLOAT, IntPtr.Zero );
                    break;
                case Format.Float3:
                    Gl.glTexImage2D( Gl.GL_TEXTURE_2D, 0, Gl.GL_RGB32F_ARB, Width, Height, 0, Gl.GL_RGB, Gl.GL_FLOAT, IntPtr.Zero );
                    break;
                case Format.Float:
                    Gl.glTexImage2D( Gl.GL_TEXTURE_2D, 0, Gl.GL_LUMINANCE32F_ARB, Width, Height, 0, Gl.GL_LUMINANCE, Gl.GL_FLOAT, IntPtr.Zero );
                    break;
                default:
                    throw new Exception( "Unsupported format for PBO" );
            }

            //Gl.glTexImage2D( Gl.GL_TEXTURE_2D, 0, Gl.GL_RGBA8, Width, Height, 0, Gl.GL_BGRA, Gl.GL_UNSIGNED_BYTE, IntPtr.Zero );

            Gl.glBindBuffer( Gl.GL_PIXEL_UNPACK_BUFFER, 0 );

            Gl.glEnable( Gl.GL_TEXTURE_2D );

            // Initialize offsets to pixel center sampling.

            float u = 0.5f / Width;
            float v = 0.5f / Height;

            Gl.glBegin( Gl.GL_QUADS );
            Gl.glTexCoord2f( u, v );
            Gl.glVertex2f( 0.0f, 0.0f );
            Gl.glTexCoord2f( 1.0f, v );
            Gl.glVertex2f( 1.0f, 0.0f );
            Gl.glTexCoord2f( 1.0f - u, 1.0f - v );
            Gl.glVertex2f( 1.0f, 1.0f );
            Gl.glTexCoord2f( u, 1.0f - v );
            Gl.glVertex2f( 0.0f, 1.0f );
            Gl.glEnd();

            Gl.glDisable( Gl.GL_TEXTURE_2D );
        }

        protected virtual void GLControl_MouseUp( object sender, MouseEventArgs e )
        {
            //throw new NotImplementedException();
        }

        protected virtual void GLControl_MouseDown( object sender, MouseEventArgs e )
        {
            //throw new NotImplementedException();
        }

        protected virtual void GLControl_MouseMove( object sender, MouseEventArgs e )
        {
            if ( Camera == null )
                return;

            float deltaX = ( e.X - mPrevMouse.X ) / (float)Width;
            float deltaY = ( e.Y - mPrevMouse.Y ) / (float)Height;

            deltaX *= mTimeDelta;
            deltaY *= mTimeDelta;

            float deltaT = deltaX;
            if ( Math.Abs( deltaX ) < Math.Abs( deltaY ) )
                deltaT = deltaY;

            switch ( e.Button )
            {
                case MouseButtons.Left:
                    Camera.Rotate( -deltaX, -deltaY );
                    CameraUpdate();
                    break;
                case MouseButtons.Right:
                    Camera.Dolly( -deltaT );
                    CameraUpdate();
                    break;
                case MouseButtons.Middle:
                    Camera.Pan( deltaX, -deltaY );
                    CameraUpdate();
                    break;
                default:
                    break;
            }

            mPrevMouse = e;
        }

        protected virtual void GLControl_KeyUp( object sender, KeyEventArgs e )
        {
            switch ( e.KeyCode )
            {
                case Keys.Escape:
                    Close();
                    break;
                default:
                    break;
            }
        }

        protected void GlControl_KeyDown( object sender, KeyEventArgs e )
        {
            if ( Camera == null )
                return;

            float step = mTimeDelta * .01f;
            switch ( e.KeyCode )
            {
                case Keys.W:
                    Camera.Dolly( step );
                    break;
                case Keys.S:
                    Camera.Dolly( -step );
                    break;
                case Keys.A:
                    Camera.Pan( -step, 0.0f );
                    break;
                case Keys.D:
                    Camera.Pan( step, 0.0f );
                    break;
                default:
                    break;
            }

            CameraUpdate();
        }

        protected override void OnClosing( CancelEventArgs e )
        {
            //we have to destroy the context here so that it is cleaned up before the GlControl
            Shutdown();

            base.OnClosing( e );
        }

        protected override void OnClosed( EventArgs e )
        {
            base.OnClosed( e );

            Application.Exit();
            Environment.Exit( 0 );
        }
    }
}
